import os
import numpy as np
from common.img_utils.pysnic.algorithms.snic import snic
import cv2
import math
import json
from skimage import measure
from skimage.segmentation import mark_boundaries
from tqdm import tqdm
import multiprocessing as mp
from scipy import ndimage
import EvalSPModule
from common.ml_utils.metrics import compute_pri_and_voi

import common.img_utils.fh.GraphOperator as go
from common.img_utils.fh.fh_label_map_generator import (
    generate_label_map,
    weight_from_edge,
    generate_label_image_for_array,
)
from analytics.parcel_segmentation.utils.crop_map_fill_superpixel import (
    fill_prob_map_into_sp,
)


def generate_img_boundary_binary_contour(
    img_arr, save_sep_labels=False, output_folder=None, num_multi_proc=10
):
    """
    Generate the binary contour mask and list of the raster has-value pixels
    Parameters
    ----------
    img_arr: ndarray
        ndarray image raster has 1 channel at least

    Return
    ----------
    img_boundary_contour_binary_mask: ndarray
        Boundary contour binary mask
    binary_mask: ndarray
        Binary region mask (image's h x w)
    sep_region_masks: list
        list of binary masks of individual regions

    """
    if not np.any(img_arr) > 0:
        binary_mask = np.zeros((img_arr.shape[0], img_arr.shape[1]))
        img_boundary_contour_binary_mask = np.zeros(
            (img_arr.shape[0], img_arr.shape[1])
        )
        label_im = np.zeros((img_arr.shape[0], img_arr.shape[1]))
    else:
        binary_mask = np.zeros((img_arr.shape[0], img_arr.shape[1]))
        if len(img_arr.shape) > 2:
            for i in range(img_arr.shape[2]):
                tmp_binary_arr = np.where(img_arr[:, :, i] > 0, 1, 0)
                binary_mask += tmp_binary_arr
            binary_mask = np.where(binary_mask > 0, 1, 0)
        else:
            binary_mask = np.where(img_arr > 0, 1, 0)

        # Separating multiple independent connected areas
        # in the binary_mask to every region binary mask
        label_im, nb_labels = ndimage.label(binary_mask)
        # label_im = measure.label(label_im, background=0, connectivity=1)

        # sep_region_masks = []
        # for i in range(nb_labels):
        #     # create an array which size is same as the mask but filled with
        #     # values that we get from the label_im.
        #     # If there are three masks, then the pixels are labeled
        #     # as 1, 2 and 3.

        #     print(f"The current isolated labels index is {i}")

        #     # List way
        #     mask_compare = np.full(np.shape(label_im), i + 1)
        #     # check equality test and have the value 1 on the location of each mask
        #     separate_mask = np.equal(label_im, mask_compare).astype(np.uint8)
        #     sep_region_masks.append(separate_mask)

        img_boundary_contours, _ = cv2.findContours(
            binary_mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
        )

        img_boundary_contour_binary_mask = np.zeros(
            (img_arr.shape[0], img_arr.shape[1], 3)
        )
        img_boundary_contour_binary_mask = cv2.drawContours(
            img_boundary_contour_binary_mask,
            img_boundary_contours,
            -1,
            (255, 255, 255),
            1,
        )
        img_boundary_contour_binary_mask = (
            img_boundary_contour_binary_mask[:, :, 0] / 255
        )

    return (
        binary_mask.astype(np.uint8),
        img_boundary_contour_binary_mask.astype(np.uint8),
        label_im,
    )


def snic_generator(
    img_arr, target_number_of_segments, compactness, img_region_binary_mask
):

    """Generate SNIC superpixel on one image ndarray

    Parameters
    ----------
    img_arr : ndarray
        Input image array
    target_number_of_segments: int
        Expected number of super pixels [int] or a iterable containing seeds
    compactness : int
        Compactness factor of SNIC
    img_region_binary_mask : ndarray
        Filter the generated grid points out of the given region (binary mask)
    Return
    ----------
    seg_label_array : ndarray
        Label map of superpixel
    actual_number_of_segments: int
        Actual number of superpixels
    """
    img_arr_lst = img_arr.tolist()

    segmentation, _, centroids = snic(
        img_arr_lst, target_number_of_segments, compactness, img_region_binary_mask
    )
    actual_number_of_segments = len(centroids)

    # Convert segmentation to local boundaries
    # The labels of seg_label_array start from -1 (nodata or background),
    # so let them start from 0. If there is no nodata or background (0s), the starting label index will be 1
    seg_label_array = np.array(segmentation)
    seg_label_array += 1

    if not np.any(seg_label_array) and np.any(img_region_binary_mask):
        # If SNIC has no outputs (all are 0s, background),
        # but the patch has raster values, mask all pixels are 1s (foreground).
        seg_label_array += img_region_binary_mask

    return seg_label_array, actual_number_of_segments


def fh_generator(
    img,
    sigma,
    k,
    min_size,
    img_type="multi-all",
    chan_num=7,
    save_results=False,
    output_folder=None,
    img_basename=None,
    patch_w_idx=None,
    patch_h_idx=None,
    patch_real_w=None,
    patch_real_h=None,
    current_count=None,
    crop_trinary_patch=None,
    save_fh_label_npy=True,
    save_img_boundaries=True,
    save_original_img=False,
    save_original_npy=False,
    save_fh_img=False,
    save_crop_prob_mask=False,
    save_overwrite=True,
    gt_binary_patch=None,
    save_metric=False,
    downscale=255.0,
    rescale=4000,
    gaussian_kernel_size=3,
):
    """Generate FH superpixel for a given image.
    Compared with the compiled version in the OpenCV,
    this FH superpixel version is able to support multi-channel image input.
    It also support multi-temporal images which are concatenated along channel dimension.
    Reference: "Felzenszwalb PF, Huttenlocher DP. Efficient graph-based image segmentation.
    In ternational journal of computer vision. 2004 Sep;59(2):167-81."

    Parameters
    ----------
    img : ndarray
        an [h x w x (cT)] image input, which should a concatenated ndarray input from the s2 and/or s1 source.
        'c' means # of channels of each image, 'T' means # of different temporal images.
        If there are multiple temporal raster sources, they should be concatenated along the channel 'c' dimension.
        For example, if a series of 7-channel rasters are obtained in 4 different time periods,
        the dimension of the accepted format is h x w x 28.
        The default value range is 0~255 float/int.
    sigma : float
        a hyper-parameter of Gaussian kernel with sigma to control the smoothness
    k : int
        k effectively sets a scale of observation,
        in that a larger k causes a preference for larger superpixel components
    min_size : int
        any segmentation area smaller than the minimum size will be merged to the neighbor large area
    img_type : str, optional
        define the image source type, in order to perform the appropriate image processing, by default 'multi-all'
    chan_num : int, optional
        # of channel of each image, by default 7, according to Chen's typical channel settings.
        It is the 'c' in [h x w x (cT)].
        When 'img_type' is 'multi-all', this variable is not used
    save_results : str, optional
        save_results is on or not, by default False.
        The filename format is --
        output_folder, img_basename + f"_w{patch_w_idx}_h{patch_h_idx}_fh_{sigma}_{k}_{min_size}"
    output_folder : int, optional
        for file saving, if save_results, by default None
    img_basename : float, optional
        for file saving, if save_results, by default None
    patch_w_idx : int, optional
        for file saving, the upper left x-axis value of the image patch, if save_results, by default None
    patch_h_idx : int, optional
        for file saving, the upper left y-axis value of the image patch, if save_results, by default None
    patch_real_w : int, optional
        for file saving, the updated patch width, by default None
    patch_real_h : int, optional
        for file saving, the updated patch height, by default None
    current_count : int, optional
        for file saving, the current count, by default None
    crop_trinary_patch: ndarray, optional
        for crop probability map guided superpixel generation
    save_fh_label_npy : bool, optional
        save the fh result in the ndarray format, if save_results, by default True
    save_img_boundaries : bool, optional
        save the boundary info of the image patch, if save_results, by default True
    save_original_img : bool, optional
        save the original image, if save_results, by default False
    save_fh_img : bool, optional
        save the fh result overlay on the image, if save_results, by default False
    save_crop_prob_mask : bool, optional
        save the crop probability map binary mask, if save_results, by default False
    save_overwrite : bool, optional
        whether overwrite the existing file, if save_results, by default True
    gt_binary_patch : ndarray, optional
        ground truth binary mask, by default None
    save_metric : bool, optional
        save the metric results, if save_metric, by default False
    downscale : float, optional
        downscale denominator of image color range,
        by default 255.0 for default input image which has [0~255] float/int.
        If the input image has [0~1], this 'downscale' value should be set as 1.0
    rescale : int, optional
        rescale multiplier of image color range, by default 4000.
        The alteration of this value is not recommended, because FH parameters are determined by this setting.
    gaussian_kernel_size : int, optional
        gaussian kernel size, by default 3.
        The alteration of this value is not recommended, because FH parameters are determined by this setting.

    Returns
    -------
    save_label_map_relabeled : ndarray
        a FH superpixel segmentation label map

    Raises
    ------
    ValueError
        the img_type must be the one from the given options
    """

    if not np.any(img) > 0:
        save_label_map_relabeled = np.zeros((img.shape[0], img.shape[1])).astype(
            np.uint8
        )
        prob_map_binary_mask = np.zeros((img.shape[0], img.shape[1])).astype(np.uint8)
        metric_dict_patch = None
    else:
        float_img = np.asarray(img, dtype=float)

        # FH paras are severely influenced image color scale,
        # the current hand-tuned paras are based on the image which have values of 0~4000.
        float_img = float_img / downscale * rescale

        gaussian_kernel = (gaussian_kernel_size, gaussian_kernel_size)
        gaussian_img = cv2.GaussianBlur(float_img, gaussian_kernel, sigma)

        if img_type == "multi-3":
            # image type RGB multiple channels
            all_chans = cv2.split(gaussian_img)
            num_img = int(len(all_chans) / chan_num)
            all_chans_3 = []
            for i in range(num_img):
                all_chans_3.extend(all_chans[i * chan_num + 0 : i * chan_num + 3])
            smooth_img = tuple(all_chans_3)
        elif img_type == "multi-4":
            # image type RGBN multiple channels
            all_chans = cv2.split(gaussian_img)
            num_img = int(len(all_chans) / chan_num)
            all_chans_4 = []
            for i in range(num_img):
                all_chans_4.extend(all_chans[i * chan_num + 0 : i * chan_num + 4])
            smooth_img = tuple(all_chans_4)
        elif img_type == "multi-all":
            # Sentinel-2 image type multiple channels
            all_chans = cv2.split(gaussian_img)
            smooth_img = tuple(all_chans)
        else:
            raise ValueError(
                "img_type for FH superpixel must be 'multi-3', 'multi-4', 'multi-all'..."
            )

        height, width, _ = img.shape
        graph = go.build_graph(smooth_img, width, height)

        sorted_graph = sorted(graph, key=weight_from_edge)

        ufset = go.segment_graph(sorted_graph, width * height, k)
        ufset = go.remove_small_component(ufset, sorted_graph, min_size)

        save_label_map = generate_label_map(ufset, width, height)

        # In this version of FH, connectivity must be 2
        save_label_map_relabeled = measure.label(
            save_label_map, background=0, connectivity=2
        )

        if crop_trinary_patch is not None:
            prob_map_binary_mask = fill_prob_map_into_sp(
                crop_trinary_patch,
                save_label_map_relabeled,
                patch_w_idx,
                patch_h_idx,
            )

            # TODO: overlay the sp results lead to incorrect stretched polygons, so have commented the following lines
            # save_label_map_updated = save_label_map_relabeled * prob_map_binary_mask
            # save_label_map_relabeled = measure.label(
            #     save_label_map_updated, background=0, connectivity=1
            # )

        else:
            prob_map_binary_mask = np.zeros((img.shape[0], img.shape[1])).astype(
                np.uint8
            )

        if gt_binary_patch is not None:
            gt_label_map_remap = measure.label(
                gt_binary_patch, background=0, connectivity=1
            ).astype(np.uint32)
            save_label_map_relabeled = save_label_map_relabeled.astype(np.uint32)

            # Calculate the eval metrics
            if np.any(gt_label_map_remap):
                # If ground truth has values, calculate the eval metrics
                label_list = save_label_map_relabeled.flatten().tolist()
                gtseg_list = gt_label_map_remap.flatten().tolist()

                # Achievable Segmentation Accuracy (ASA)
                # Bergh Underseg Error (UE)
                asa = EvalSPModule.computeASA(label_list, gtseg_list, 0)
                ue = 1 - asa

                # PRI and VoI
                curRI, curVoI = compute_pri_and_voi(
                    save_label_map_relabeled, gt_label_map_remap
                )

                # Segmentation Counts
                perc_1 = 0.1
                (
                    nr_underseg_1,
                    nr_one_to_one_1,
                    nr_overseg_1,
                    nr_gt_1,
                ) = EvalSPModule.computeSegCatCounts(label_list, gtseg_list, perc_1, 0)

                perc_2 = 0.2
                (
                    nr_underseg_2,
                    nr_one_to_one_2,
                    nr_overseg_2,
                    nr_gt_2,
                ) = EvalSPModule.computeSegCatCounts(label_list, gtseg_list, perc_2, 0)

                perc_3 = 0.3
                (
                    nr_underseg_3,
                    nr_one_to_one_3,
                    nr_overseg_3,
                    nr_gt_3,
                ) = EvalSPModule.computeSegCatCounts(label_list, gtseg_list, perc_3, 0)

                # Boundary Recall (BR)
                rad_1 = 1
                (
                    br_gt_1,
                    br_sp_1,
                    average_error_pixel_gt_1,
                    average_error_pixel_sp_1,
                ) = EvalSPModule.computeBR(label_list, gtseg_list, height, width, rad_1)

                rad_2 = 2
                (
                    br_gt_2,
                    br_sp_2,
                    average_error_pixel_gt_2,
                    average_error_pixel_sp_2,
                ) = EvalSPModule.computeBR(label_list, gtseg_list, height, width, rad_2)

                rad_3 = 3
                (
                    br_gt_3,
                    br_sp_3,
                    average_error_pixel_gt_3,
                    average_error_pixel_sp_3,
                ) = EvalSPModule.computeBR(label_list, gtseg_list, height, width, rad_3)

                # Construct eval metric into a dict for json
                metric_dict_patch = dict(
                    patch_width=patch_real_w,
                    patch_height=patch_real_h,
                    top_corner_y=patch_h_idx,
                    top_corner_x=patch_w_idx,
                    current_count=current_count,
                    asa=asa,
                    ue=ue,
                    pri=curRI,
                    voi=curVoI,
                    perc_1=perc_1,
                    number_gt_labels_1=nr_gt_1,
                    nr_underseg_one2one_overseg_1=[
                        nr_underseg_1,
                        nr_one_to_one_1,
                        nr_overseg_1,
                    ],
                    perc_2=perc_2,
                    number_gt_labels_2=nr_gt_2,
                    nr_underseg_one2one_overseg_2=[
                        nr_underseg_2,
                        nr_one_to_one_2,
                        nr_overseg_2,
                    ],
                    perc_3=perc_3,
                    number_gt_labels_3=nr_gt_3,
                    nr_underseg_one2one_overseg_3=[
                        nr_underseg_3,
                        nr_one_to_one_3,
                        nr_overseg_3,
                    ],
                    rad_1=rad_1,
                    br_gt_1=br_gt_1,
                    br_sp_1=br_sp_1,
                    average_error_pixel_gt_1=average_error_pixel_gt_1,
                    average_error_pixel_sp_1=average_error_pixel_sp_1,
                    rad_2=rad_2,
                    br_gt_2=br_gt_2,
                    br_sp_2=br_sp_2,
                    average_error_pixel_gt_2=average_error_pixel_gt_2,
                    average_error_pixel_sp_2=average_error_pixel_sp_2,
                    rad_3=rad_3,
                    br_gt_3=br_gt_3,
                    br_sp_3=br_sp_3,
                    average_error_pixel_gt_3=average_error_pixel_gt_3,
                    average_error_pixel_sp_3=average_error_pixel_sp_3,
                )
            else:
                metric_dict_patch = None
        else:
            metric_dict_patch = None

    if save_results:
        save_fh_results(
            img,
            save_label_map_relabeled,
            output_folder,
            img_basename,
            patch_w_idx,
            patch_h_idx,
            sigma,
            k,
            min_size,
            crop_prob_mask=prob_map_binary_mask,
            save_fh_label_npy=save_fh_label_npy,
            save_img_boundaries=save_img_boundaries,
            save_original_img=save_original_img,
            save_original_npy=save_original_npy,
            save_fh_img=save_fh_img,
            save_crop_prob_mask=save_crop_prob_mask,
            save_overwrite=save_overwrite,
            save_metric=save_metric,
            metric_dict=metric_dict_patch,
        )

    return save_label_map_relabeled


def save_fh_results(  # noqa: C901
    img,
    label_map,
    output_folder,
    img_basename,
    patch_w_idx,
    patch_h_idx,
    fh_sigma,
    fh_k,
    fh_min_size,
    crop_prob_mask=None,
    save_fh_label_npy=True,
    save_img_boundaries=True,
    save_original_img=False,
    save_original_npy=False,
    save_fh_img=False,
    save_crop_prob_mask=False,
    save_overwrite=False,
    save_metric=False,
    metric_dict=None,
):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    if save_fh_label_npy:
        save_filepath = os.path.join(
            output_folder,
            img_basename
            + f"_w{patch_w_idx}_h{patch_h_idx}_lbl_fh_{fh_sigma}_{fh_k}_{fh_min_size}.npy",
        )
        if os.path.exists(save_filepath):
            # The file has been generated before
            if save_overwrite:
                np.save(save_filepath, label_map)
        else:
            np.save(save_filepath, label_map)

    if save_img_boundaries:
        (
            img_region_binary_mask,
            img_boundary_binary_mask,
            _,
        ) = generate_img_boundary_binary_contour(img)

        save_filepath_region_mask = os.path.join(
            output_folder,
            img_basename + f"_w{patch_w_idx}_h{patch_h_idx}_region_mask.npy",
        )
        save_filepath_boundary_mask = os.path.join(
            output_folder,
            img_basename + f"_w{patch_w_idx}_h{patch_h_idx}_boundary_mask.npy",
        )

        if os.path.exists(save_filepath):
            # The file has been generated before
            if save_overwrite:
                np.save(save_filepath_region_mask, img_region_binary_mask)
                np.save(save_filepath_boundary_mask, img_boundary_binary_mask)
        else:
            np.save(save_filepath_region_mask, img_region_binary_mask)
            np.save(save_filepath_boundary_mask, img_boundary_binary_mask)

    if save_original_npy:
        save_filepath_img_npy = os.path.join(
            output_folder,
            img_basename + f"_w{patch_w_idx}_h{patch_h_idx}.npy",
        )
        np.save(save_filepath_img_npy, img)

    if save_original_img or save_fh_img:
        img_for_visual = img[:, :, [2, 1, 0]].astype(np.uint8)

    if save_original_img:
        save_filepath = os.path.join(
            output_folder, img_basename + f"_w{patch_w_idx}_h{patch_h_idx}.png"
        )
        if os.path.exists(save_filepath):
            # The file has been generated before
            if save_overwrite:
                cv2.imwrite(save_filepath, img_for_visual)
        else:
            cv2.imwrite(save_filepath, img_for_visual)

    if save_fh_img:
        save_filepath = os.path.join(
            output_folder,
            img_basename
            + f"_w{patch_w_idx}_h{patch_h_idx}_fh_{fh_sigma}_{fh_k}_{fh_min_size}.png",
        )
        if os.path.exists(save_filepath):
            # The file has been generated before
            if save_overwrite:
                cv2.imwrite(
                    save_filepath, mark_boundaries(img_for_visual, label_map) * 255
                )
        else:
            cv2.imwrite(save_filepath, mark_boundaries(img_for_visual, label_map) * 255)

    if save_crop_prob_mask and crop_prob_mask is not None:
        save_filepath = os.path.join(
            output_folder,
            img_basename + f"_w{patch_w_idx}_h{patch_h_idx}_crop_mask.png",
        )
        if os.path.exists(save_filepath):
            # The file has been generated before
            if save_overwrite:
                cv2.imwrite(
                    save_filepath,
                    (crop_prob_mask * 255).astype(np.uint8),
                )
        else:
            cv2.imwrite(
                save_filepath,
                (crop_prob_mask * 255).astype(np.uint8),
            )

    if save_metric and metric_dict is not None:
        save_filepath = os.path.join(
            output_folder,
            img_basename + f"_w{patch_w_idx}_h{patch_h_idx}_eval.json",
        )
        with open(save_filepath, "w") as f:
            json.dump(metric_dict, f, indent=2)


def crop_img_and_fh_on_patch(  # noqa: C901
    img_arr,
    output_folder,
    img_basename,
    patch_w,
    patch_h,
    fh_sigma,
    fh_k,
    fh_min_size,
    fh_img_type="multi-4",
    chan_num=7,
    crop_trinary_mask: np.ndarray = None,
    save_results=True,
    save_fh_label_npy=True,
    save_img_boundaries=True,
    save_original_img=False,
    save_original_npy=False,
    save_fh_img=False,
    save_crop_prob_mask=False,
    save_overwrite=True,
    gt_binary_mask: np.ndarray = None,
    save_metric=False,
    downscale=255.0,
    num_multi_proc=10,
):
    """(Multi-Processing) Crop a given multi-temporal multi-channel large image into patches
    and generate FH superpixel on each patch.
    Compared with the compiled version in the OpenCV,
    this FH superpixel version is able to support multi-channel image input.
    It also support multi-temporal images which are concatenated along channel dimension.
    Reference: "Felzenszwalb PF, Huttenlocher DP. Efficient graph-based image segmentation.
    In ternational journal of computer vision. 2004 Sep;59(2):167-81."

    Parameters
    ----------
    img_arr : ndarray
        an [h x w x (cT)] image input, which should a concatenated ndarray input from the s2 and/or s1 source.
        'c' means # of channels of each image, 'T' means # of different temporal images.
        If there are multiple temporal raster sources, they should be concatenated along the channel 'c' dimension.
        For example, if a series of 7-channel rasters are obtained in 4 different time periods,
        the dimension of the accepted format is h x w x 28.
        The default value range is 0~255 float/int.
    output_folder : int, optional
        for file saving
    img_basename : float, optional
        for file saving
    patch_w : int, optional
        width of the image patch
    patch_h : int, optional
        height of the image patch
    fh_sigma : float
        a hyper-parameter of Gaussian kernel with sigma to control the smoothness in the FH
    fh_k : int
        k effectively sets a scale of observation in the FH,
        in that a larger k causes a preference for larger superpixel components
    fh_min_size : int
        any segmentation area smaller than the minimum size will be merged to the neighbor large area in the FH
    fh_img_type : str, optional
        define the image source type in the FH,
        in order to perform the appropriate image processing, by default 'multi-all'
    chan_num : int, optional
        # of channel of each image, by default 7, according to Chen's typical channel settings.
        It is the 'c' in [h x w x (cT)].
        When 'fh_img_type' is 'multi-all', this variable is not used
    crop_prob_map: ndarray, optional
        crop probability map input ([0-201] integer ndarray)
    crop_binary_thres: float, optional
        threshold for crop probability map into a trinary mask ([0, 1, 2] integer ndarray).
        0 is nodata, 1 is non-crop, 2 is crop
    save_results : str, optional
        save_results is on or not, by default False
        The filename format is --
        output_folder, img_basename + f"_w{patch_w}_h{patch_h}_fh_{fh_sigma}_{fh_k}_{fh_min_size}"
    save_fh_label_npy : bool, optional
        save the fh result in the ndarray format, if save_results, by default True
    save_original_img : bool, optional
        save the original image, if save_results, by default False
    save_fh_img : bool, optional
        save the fh result overlay on the image, if save_results, by default False
    save_crop_prob_mask : bool, optional
        save the crop probability map binary mask, if save_results, by default False
    save_overwrite : bool, optional
        whether overwrite the existing file, if save_results, by default True
    gt_binary_mask : ndarray, optional
        ground truth binary map binary mask, by default None
    save_metric : bool, optional
        save the metric results, if save_metric, by default False
    downscale : float, optional
        downscale denominator of image color range,
        by default 255.0 for default input image which has [0~255] float/int.
        If the input image has [0~1], this 'downscale' value should be set as 1.0
    num_multi_proc: int, optional
        # of multiple processors for multiple processing, by default 10.

    Returns
    -------
    save_label_map_relabeled : ndarray
        a FH superpixel segmentation label map
    """

    print("Start cropping images...")
    img_h, img_w, _ = img_arr.shape
    num_row = math.ceil(img_h / patch_h)
    num_col = math.ceil(img_w / patch_w)

    current_count = 0
    pool = mp.Pool(num_multi_proc)

    if crop_trinary_mask is None:
        crop_tri_mask_patch = None

    if gt_binary_mask is None:
        gt_binary_mask_patch = None

    for col in range(num_col):
        for row in range(num_row):

            current_count += 1

            if col == num_col - 1 and row == num_row - 1:
                crop_img = img_arr[row * patch_h :, col * patch_w :]
                if crop_trinary_mask is not None:
                    crop_tri_mask_patch = crop_trinary_mask[
                        row * patch_h :, col * patch_w :
                    ]
                if gt_binary_mask is not None:
                    gt_binary_mask_patch = gt_binary_mask[
                        row * patch_h :, col * patch_w :
                    ]
                patch_real_h = img_h - row * patch_h  # for metric
                patch_real_w = img_w - col * patch_w  # for metric
            elif col == num_col - 1:
                crop_img = img_arr[row * patch_h : (row + 1) * patch_h, col * patch_w :]
                if crop_trinary_mask is not None:
                    crop_tri_mask_patch = crop_trinary_mask[
                        row * patch_h : (row + 1) * patch_h, col * patch_w :
                    ]
                if gt_binary_mask is not None:
                    gt_binary_mask_patch = gt_binary_mask[
                        row * patch_h : (row + 1) * patch_h, col * patch_w :
                    ]
                patch_real_h = patch_h
                patch_real_w = img_w - col * patch_w
            elif row == num_row - 1:
                crop_img = img_arr[row * patch_h :, col * patch_w : (col + 1) * patch_w]
                if crop_trinary_mask is not None:
                    crop_tri_mask_patch = crop_trinary_mask[
                        row * patch_h :, col * patch_w : (col + 1) * patch_w
                    ]
                if gt_binary_mask is not None:
                    gt_binary_mask_patch = gt_binary_mask[
                        row * patch_h :, col * patch_w : (col + 1) * patch_w
                    ]
                patch_real_h = img_h - row * patch_h
                patch_real_w = patch_w
            else:
                crop_img = img_arr[
                    row * patch_h : (row + 1) * patch_h,
                    col * patch_w : (col + 1) * patch_w,
                ]
                if crop_trinary_mask is not None:
                    crop_tri_mask_patch = crop_trinary_mask[
                        row * patch_h : (row + 1) * patch_h,
                        col * patch_w : (col + 1) * patch_w,
                    ]
                if gt_binary_mask is not None:
                    gt_binary_mask_patch = gt_binary_mask[
                        row * patch_h : (row + 1) * patch_h,
                        col * patch_w : (col + 1) * patch_w,
                    ]
                patch_real_h = patch_h
                patch_real_w = patch_w

            top_corner_x = col * patch_w
            top_corner_y = row * patch_h

            # Debug #
            # if top_corner_x < 2048:
            #     print(patch_h * row)
            #     continue
            # if top_corner_y > 0:
            #     print(patch_w * col)
            #     continue

            pool.apply_async(
                fh_generator,
                args=[
                    crop_img,
                    fh_sigma,
                    fh_k,
                    fh_min_size,
                    fh_img_type,
                    chan_num,
                    save_results,
                    output_folder,
                    img_basename,
                    top_corner_x,
                    top_corner_y,
                    patch_real_w,
                    patch_real_h,
                    current_count,
                    crop_tri_mask_patch,
                    save_fh_label_npy,
                    save_img_boundaries,
                    save_original_img,
                    save_original_npy,
                    save_fh_img,
                    save_crop_prob_mask,
                    save_overwrite,
                    gt_binary_mask_patch,
                    save_metric,
                    downscale,
                ],
            )

    pool.close()
    pool.join()


def mosaic_fh_patches_to_img(
    load_folder,
    original_img_arr,
    input_w=512,
    input_h=512,
    load_img_type=".npy",
    img_region_type="tile",
    output_npy_path=None,
    visualize_is_on=False,
    save_npy=True,
):
    """Mosaic the label map patches into a whole ndarray.

    Parameters
    ----------
    load_folder : string
        The files in a load folder format is default as --
        load_folder, img_basename + f"_w{patch_w}_h{patch_h}_fh_{fh_sigma}_{fh_k}_{fh_min_size}"
        If the format is not that format, it is unable to be properly mosaicked at this point.
    output_npy_path : string
        The save path for the mosaicked label map ndarray.
    original_img_arr : ndarray
        The original image in the ndarray format, [0~255] float/int.
    input_w : int, optional
        width of each patch ndarray, by default 512
    input_h : int, optional
        height of each patch ndarray, by default 512
    load_img_type : str, optional
        the file type will be loaded in the load folder, by default '.npy'
    img_region_type : str, optional
        the original image region is 'region' (has nodata boundary) or 'tile', by default 'tile'.
        If 'region', the label indexes in the nodata region will be removed.
    visualize_is_on : bool, optional
        visualization is on or not for the mosaicked ndarray, by default False
    save_npy : bool, optional
        save_npy is on or not for the mosaicked ndarray, by default True

    Returns
    -------
    save_label_map_relabeled : ndarray
        a FH superpixel segmentation label map
    """

    print("Start to mosaic ...")
    fh_labels_filenames = [
        f
        for f in os.listdir(load_folder)
        if f.endswith(load_img_type) and "gt" not in f and "lbl" in f
    ]

    out_h, out_w, _ = original_img_arr.shape

    img_out_label = np.zeros((out_h, out_w)).astype(np.uint32)
    if visualize_is_on:
        img_out = np.zeros((out_h, out_w, 3)).astype(np.uint8)

    cur_label_end_idx = None

    # SP Patch Merge
    for i in tqdm(range(len(fh_labels_filenames))):
        fh_labels = np.load(os.path.join(load_folder, fh_labels_filenames[i]))
        if cur_label_end_idx is None:
            cur_label_end_idx = np.max(fh_labels)
        else:
            cur_label_end_idx = cur_label_end_idx + np.max(fh_labels) + 1

        fh_labels_updated = fh_labels.astype(np.uint32)
        fh_labels_updated += cur_label_end_idx

        tmp_w = [
            i
            for i in os.path.splitext(fh_labels_filenames[i])[0].split("_")
            if i.startswith("w")
        ][0][1:]

        tmp_h = [
            i
            for i in os.path.splitext(fh_labels_filenames[i])[0].split("_")
            if i.startswith("h")
        ][0][1:]

        w, h = eval(tmp_w), eval(tmp_h)

        if out_h - h < input_h and out_w - w < input_w:
            offset_h = out_h - h
            offset_w = out_w - w
        elif out_h - h < input_h:
            offset_h = out_h - h
            offset_w = input_w
        elif out_w - w < input_w:
            offset_h = input_h
            offset_w = out_w - w
        else:
            offset_h = input_h
            offset_w = input_w

        img_out_label[h : h + offset_h, w : w + offset_w] = fh_labels_updated
        if visualize_is_on:
            generated_colored_label_map = generate_label_image_for_array(fh_labels)
            img_out[h : h + offset_h, w : w + offset_w] = generated_colored_label_map

    if img_region_type == "region":
        img_region_binary_mask, _, _ = generate_img_boundary_binary_contour(
            original_img_arr
        )

        img_out_label += 1
        img_out_label = img_out_label * img_region_binary_mask
        img_out_label = measure.label(img_out_label, background=0, connectivity=1)

    if save_npy and output_npy_path is not None:
        print(f"Save the mosaicked npy... {output_npy_path}")
        np.save(output_npy_path, img_out_label)

    if visualize_is_on:
        output_img_path = output_npy_path.replace(".npy", ".png")
        print(f"Save the mosaicked png... {output_img_path}")
        cv2.imwrite(output_img_path, img_out)

    return img_out_label
